// Copyright 2014 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "net/spdy/http2_write_scheduler.h"

#include "net/spdy/spdy_test_utils.h"
#include "net/test/gtest_util.h"
#include "testing/gmock/include/gmock/gmock.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace net {

using ::testing::AssertionFailure;
using ::testing::AssertionResult;
using ::testing::AssertionSuccess;
using ::testing::ElementsAre;
using ::testing::IsEmpty;
using ::testing::UnorderedElementsAre;

namespace test {

    template <typename StreamIdType>
    class Http2PriorityWriteSchedulerPeer {
    public:
        explicit Http2PriorityWriteSchedulerPeer(
            Http2PriorityWriteScheduler<StreamIdType>* scheduler)
            : scheduler_(scheduler)
        {
        }

        int TotalChildWeights(StreamIdType stream_id) const
        {
            return scheduler_->FindStream(stream_id)->total_child_weights;
        }

        bool ValidateInvariants() const
        {
            return scheduler_->ValidateInvariantsForTests();
        }

    private:
        Http2PriorityWriteScheduler<StreamIdType>* scheduler_;
    };

    class Http2PriorityWriteSchedulerTest : public ::testing::Test {
    protected:
        using SpdyStreamId = uint32_t;

        Http2PriorityWriteSchedulerTest()
            : peer_(&scheduler_)
        {
        }

        Http2PriorityWriteScheduler<SpdyStreamId> scheduler_;
        Http2PriorityWriteSchedulerPeer<SpdyStreamId> peer_;
    };

    TEST_F(Http2PriorityWriteSchedulerTest, RegisterAndUnregisterStreams)
    {
        EXPECT_EQ(1, scheduler_.num_streams());
        EXPECT_TRUE(scheduler_.StreamRegistered(0));
        EXPECT_FALSE(scheduler_.StreamRegistered(1));

        scheduler_.RegisterStream(1, SpdyStreamPrecedence(0, 100, false));
        EXPECT_EQ(2, scheduler_.num_streams());
        ASSERT_TRUE(scheduler_.StreamRegistered(1));
        EXPECT_EQ(100, scheduler_.GetStreamPrecedence(1).weight());
        EXPECT_FALSE(scheduler_.StreamRegistered(5));
        EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(1));

        scheduler_.RegisterStream(5, SpdyStreamPrecedence(0, 50, false));
        // Should not be able to add a stream with an id that already exists.
        EXPECT_SPDY_BUG(
            scheduler_.RegisterStream(5, SpdyStreamPrecedence(1, 50, false)),
            "Stream 5 already registered");
        EXPECT_EQ(3, scheduler_.num_streams());
        EXPECT_TRUE(scheduler_.StreamRegistered(1));
        ASSERT_TRUE(scheduler_.StreamRegistered(5));
        EXPECT_EQ(50, scheduler_.GetStreamPrecedence(5).weight());
        EXPECT_FALSE(scheduler_.StreamRegistered(13));

        scheduler_.RegisterStream(13, SpdyStreamPrecedence(5, 130, true));
        EXPECT_EQ(4, scheduler_.num_streams());
        EXPECT_TRUE(scheduler_.StreamRegistered(1));
        EXPECT_TRUE(scheduler_.StreamRegistered(5));
        ASSERT_TRUE(scheduler_.StreamRegistered(13));
        EXPECT_EQ(130, scheduler_.GetStreamPrecedence(13).weight());
        EXPECT_EQ(5u, scheduler_.GetStreamPrecedence(13).parent_id());

        scheduler_.UnregisterStream(5);
        // Cannot remove a stream that has already been removed.
        EXPECT_SPDY_BUG(scheduler_.UnregisterStream(5), "Stream 5 not registered");
        EXPECT_EQ(3, scheduler_.num_streams());
        EXPECT_TRUE(scheduler_.StreamRegistered(1));
        EXPECT_FALSE(scheduler_.StreamRegistered(5));
        EXPECT_TRUE(scheduler_.StreamRegistered(13));
        EXPECT_EQ(kHttp2RootStreamId, scheduler_.GetStreamPrecedence(13).parent_id());

        // The parent stream 19 doesn't exist, so this should use 0 as parent stream:
        EXPECT_SPDY_BUG(
            scheduler_.RegisterStream(7, SpdyStreamPrecedence(19, 70, false)),
            "Parent stream 19 not registered");
        EXPECT_TRUE(scheduler_.StreamRegistered(7));
        EXPECT_EQ(0u, scheduler_.GetStreamPrecedence(7).parent_id());
        // Now stream 7 already exists, so this should fail:
        EXPECT_SPDY_BUG(
            scheduler_.RegisterStream(7, SpdyStreamPrecedence(1, 70, false)),
            "Stream 7 already registered");
        // Try adding a second child to stream 13:
        scheduler_.RegisterStream(17, SpdyStreamPrecedence(13, 170, false));

        scheduler_.UpdateStreamPrecedence(17, SpdyStreamPrecedence(13, 150, false));
        EXPECT_EQ(150, scheduler_.GetStreamPrecedence(17).weight());

        ASSERT_TRUE(peer_.ValidateInvariants());
    }

    TEST_F(Http2PriorityWriteSchedulerTest, RegisterStreamWithSpdy3Priority)
    {
        EXPECT_FALSE(scheduler_.StreamRegistered(1));
        EXPECT_SPDY_BUG(scheduler_.RegisterStream(1, SpdyStreamPrecedence(3)),
            "Expected HTTP/2 stream dependency");
        EXPECT_EQ(0u, scheduler_.NumReadyStreams());
        EXPECT_TRUE(scheduler_.StreamRegistered(1));
        EXPECT_EQ(3, scheduler_.GetStreamPrecedence(1).spdy3_priority());
        EXPECT_EQ(147, scheduler_.GetStreamPrecedence(1).weight());
        EXPECT_EQ(kHttp2RootStreamId, scheduler_.GetStreamPrecedence(1).parent_id());
        EXPECT_THAT(scheduler_.GetStreamChildren(1), IsEmpty());

        EXPECT_SPDY_BUG(scheduler_.RegisterStream(1, SpdyStreamPrecedence(4)),
            "Stream 1 already registered");
        EXPECT_EQ(3, scheduler_.GetStreamPrecedence(1).spdy3_priority());
    }

    TEST_F(Http2PriorityWriteSchedulerTest, GetStreamWeight)
    {
        EXPECT_SPDY_BUG(EXPECT_EQ(kHttp2MinStreamWeight,
                            scheduler_.GetStreamPrecedence(3).weight()),
            "Stream 3 not registered");
        scheduler_.RegisterStream(3, SpdyStreamPrecedence(0, 130, true));
        EXPECT_EQ(130, scheduler_.GetStreamPrecedence(3).weight());
        scheduler_.UpdateStreamPrecedence(3, SpdyStreamPrecedence(0, 50, true));
        EXPECT_EQ(50, scheduler_.GetStreamPrecedence(3).weight());
        scheduler_.UnregisterStream(3);
        EXPECT_SPDY_BUG(EXPECT_EQ(kHttp2MinStreamWeight,
                            scheduler_.GetStreamPrecedence(3).weight()),
            "Stream 3 not registered");
    }

    TEST_F(Http2PriorityWriteSchedulerTest, GetStreamPriority)
    {
        EXPECT_SPDY_BUG(EXPECT_EQ(kV3LowestPriority,
                            scheduler_.GetStreamPrecedence(3).spdy3_priority()),
            "Stream 3 not registered");
        scheduler_.RegisterStream(3, SpdyStreamPrecedence(0, 130, true));
        EXPECT_EQ(3, scheduler_.GetStreamPrecedence(3).spdy3_priority());
        scheduler_.UpdateStreamPrecedence(3, SpdyStreamPrecedence(0, 50, true));
        EXPECT_EQ(5, scheduler_.GetStreamPrecedence(3).spdy3_priority());
        scheduler_.UnregisterStream(3);
        EXPECT_SPDY_BUG(EXPECT_EQ(kV3LowestPriority,
                            scheduler_.GetStreamPrecedence(3).spdy3_priority()),
            "Stream 3 not registered");
    }

    TEST_F(Http2PriorityWriteSchedulerTest, GetStreamParent)
    {
        EXPECT_SPDY_BUG(EXPECT_EQ(kHttp2RootStreamId,
                            scheduler_.GetStreamPrecedence(3).parent_id()),
            "Stream 3 not registered");
        scheduler_.RegisterStream(2, SpdyStreamPrecedence(0, 20, false));
        scheduler_.RegisterStream(3, SpdyStreamPrecedence(2, 30, false));
        EXPECT_EQ(2u, scheduler_.GetStreamPrecedence(3).parent_id());
        scheduler_.UnregisterStream(3);
        EXPECT_SPDY_BUG(EXPECT_EQ(kHttp2RootStreamId,
                            scheduler_.GetStreamPrecedence(3).parent_id()),
            "Stream 3 not registered");
    }

    TEST_F(Http2PriorityWriteSchedulerTest, GetStreamChildren)
    {
        EXPECT_SPDY_BUG(EXPECT_THAT(scheduler_.GetStreamChildren(7), IsEmpty()),
            "Stream 7 not registered");
        scheduler_.RegisterStream(7, SpdyStreamPrecedence(0, 70, false));
        EXPECT_THAT(scheduler_.GetStreamChildren(7), IsEmpty());
        scheduler_.RegisterStream(9, SpdyStreamPrecedence(7, 90, false));
        scheduler_.RegisterStream(15, SpdyStreamPrecedence(7, 150, false));
        EXPECT_THAT(scheduler_.GetStreamChildren(7), UnorderedElementsAre(9, 15));
        scheduler_.UnregisterStream(7);
        EXPECT_SPDY_BUG(EXPECT_THAT(scheduler_.GetStreamChildren(7), IsEmpty()),
            "Stream 7 not registered");
    }

    TEST_F(Http2PriorityWriteSchedulerTest, UpdateStreamWeight)
    {
        EXPECT_SPDY_BUG(
            scheduler_.UpdateStreamPrecedence(0, SpdyStreamPrecedence(0, 10, false)),
            "Cannot set precedence of root stream");
        EXPECT_SPDY_BUG(
            scheduler_.UpdateStreamPrecedence(3, SpdyStreamPrecedence(0, 10, false)),
            "Stream 3 not registered");
        scheduler_.RegisterStream(3, SpdyStreamPrecedence(0, 10, false));
        scheduler_.UpdateStreamPrecedence(3, SpdyStreamPrecedence(0, 20, false));
        EXPECT_EQ(20, scheduler_.GetStreamPrecedence(3).weight());
        ASSERT_TRUE(peer_.ValidateInvariants());

        EXPECT_SPDY_BUG(
            scheduler_.UpdateStreamPrecedence(3, SpdyStreamPrecedence(0, 500, false)),
            "Invalid weight: 500");
        EXPECT_EQ(kHttp2MaxStreamWeight, scheduler_.GetStreamPrecedence(3).weight());
        EXPECT_SPDY_BUG(
            scheduler_.UpdateStreamPrecedence(3, SpdyStreamPrecedence(0, 0, false)),
            "Invalid weight: 0");
        EXPECT_EQ(kHttp2MinStreamWeight, scheduler_.GetStreamPrecedence(3).weight());
        ASSERT_TRUE(peer_.ValidateInvariants());

        scheduler_.UnregisterStream(3);
        EXPECT_SPDY_BUG(
            scheduler_.UpdateStreamPrecedence(3, SpdyStreamPrecedence(0, 10, false)),
            "Stream 3 not registered");
    }

    // Basic case of reparenting a subtree.
    TEST_F(Http2PriorityWriteSchedulerTest, UpdateStreamParentBasicNonExclusive)
    {
        /* Tree:
        0
       / \
      1   2
     / \
    3   4
   */
        scheduler_.RegisterStream(1, SpdyStreamPrecedence(0, 100, false));
        scheduler_.RegisterStream(2, SpdyStreamPrecedence(0, 100, false));
        scheduler_.RegisterStream(3, SpdyStreamPrecedence(1, 100, false));
        scheduler_.RegisterStream(4, SpdyStreamPrecedence(1, 100, false));
        scheduler_.UpdateStreamPrecedence(1, SpdyStreamPrecedence(2, 100, false));
        EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(2));
        EXPECT_THAT(scheduler_.GetStreamChildren(1), UnorderedElementsAre(3, 4));
        EXPECT_THAT(scheduler_.GetStreamChildren(2), ElementsAre(1));
        EXPECT_THAT(scheduler_.GetStreamChildren(3), IsEmpty());
        EXPECT_THAT(scheduler_.GetStreamChildren(4), IsEmpty());
        ASSERT_TRUE(peer_.ValidateInvariants());
    }

    // Basic case of reparenting a subtree.  Result here is the same as the
    // non-exclusive case.
    TEST_F(Http2PriorityWriteSchedulerTest, UpdateStreamParentBasicExclusive)
    {
        /* Tree:
        0
       / \
      1   2
     / \
    3   4
   */
        scheduler_.RegisterStream(1, SpdyStreamPrecedence(0, 100, false));
        scheduler_.RegisterStream(2, SpdyStreamPrecedence(0, 100, false));
        scheduler_.RegisterStream(3, SpdyStreamPrecedence(1, 100, false));
        scheduler_.RegisterStream(4, SpdyStreamPrecedence(1, 100, false));
        scheduler_.UpdateStreamPrecedence(1, SpdyStreamPrecedence(2, 100, true));
        EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(2));
        EXPECT_THAT(scheduler_.GetStreamChildren(1), UnorderedElementsAre(3, 4));
        EXPECT_THAT(scheduler_.GetStreamChildren(2), ElementsAre(1));
        EXPECT_THAT(scheduler_.GetStreamChildren(3), IsEmpty());
        EXPECT_THAT(scheduler_.GetStreamChildren(4), IsEmpty());
        ASSERT_TRUE(peer_.ValidateInvariants());
    }

    // We can't set the parent of a nonexistent stream, or set the parent to a
    // nonexistent stream.
    TEST_F(Http2PriorityWriteSchedulerTest, UpdateStreamParentNonexistent)
    {
        scheduler_.RegisterStream(1, SpdyStreamPrecedence(0, 100, false));
        scheduler_.RegisterStream(2, SpdyStreamPrecedence(0, 100, false));
        for (bool exclusive : { true, false }) {
            EXPECT_SPDY_BUG(scheduler_.UpdateStreamPrecedence(
                                1, SpdyStreamPrecedence(3, 100, exclusive)),
                "Parent stream 3 not registered");
            EXPECT_SPDY_BUG(scheduler_.UpdateStreamPrecedence(
                                4, SpdyStreamPrecedence(2, 100, exclusive)),
                "Stream 4 not registered");
            EXPECT_SPDY_BUG(scheduler_.UpdateStreamPrecedence(
                                3, SpdyStreamPrecedence(4, 100, exclusive)),
                "Stream 3 not registered");
            EXPECT_THAT(scheduler_.GetStreamChildren(0), UnorderedElementsAre(1, 2));
            EXPECT_THAT(scheduler_.GetStreamChildren(1), IsEmpty());
            EXPECT_THAT(scheduler_.GetStreamChildren(2), IsEmpty());
        }
        ASSERT_TRUE(peer_.ValidateInvariants());
    }

    // We should be able to add multiple children to streams.
    TEST_F(Http2PriorityWriteSchedulerTest,
        UpdateStreamParentMultipleChildrenNonExclusive)
    {
        /* Tree:
        0
       / \
      1   2
     / \   \
    3   4   5
   */
        scheduler_.RegisterStream(1, SpdyStreamPrecedence(0, 100, false));
        scheduler_.RegisterStream(2, SpdyStreamPrecedence(0, 100, false));
        scheduler_.RegisterStream(3, SpdyStreamPrecedence(1, 100, false));
        scheduler_.RegisterStream(4, SpdyStreamPrecedence(1, 100, false));
        scheduler_.RegisterStream(5, SpdyStreamPrecedence(2, 100, false));
        scheduler_.UpdateStreamPrecedence(2, SpdyStreamPrecedence(1, 100, false));
        EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(1));
        EXPECT_THAT(scheduler_.GetStreamChildren(1), UnorderedElementsAre(2, 3, 4));
        EXPECT_THAT(scheduler_.GetStreamChildren(2), ElementsAre(5));
        EXPECT_THAT(scheduler_.GetStreamChildren(3), IsEmpty());
        EXPECT_THAT(scheduler_.GetStreamChildren(4), IsEmpty());
        EXPECT_THAT(scheduler_.GetStreamChildren(5), IsEmpty());
        ASSERT_TRUE(peer_.ValidateInvariants());
    }

    TEST_F(Http2PriorityWriteSchedulerTest,
        UpdateStreamParentMultipleChildrenExclusive)
    {
        /* Tree:
        0
       / \
      1   2
     / \   \
    3   4   5
   */
        scheduler_.RegisterStream(1, SpdyStreamPrecedence(0, 100, false));
        scheduler_.RegisterStream(2, SpdyStreamPrecedence(0, 100, false));
        scheduler_.RegisterStream(3, SpdyStreamPrecedence(1, 100, false));
        scheduler_.RegisterStream(4, SpdyStreamPrecedence(1, 100, false));
        scheduler_.RegisterStream(5, SpdyStreamPrecedence(2, 100, false));
        scheduler_.UpdateStreamPrecedence(2, SpdyStreamPrecedence(1, 100, true));
        EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(1));
        EXPECT_THAT(scheduler_.GetStreamChildren(1), ElementsAre(2));
        EXPECT_THAT(scheduler_.GetStreamChildren(2), UnorderedElementsAre(3, 4, 5));
        EXPECT_THAT(scheduler_.GetStreamChildren(3), IsEmpty());
        EXPECT_THAT(scheduler_.GetStreamChildren(4), IsEmpty());
        EXPECT_THAT(scheduler_.GetStreamChildren(5), IsEmpty());
        ASSERT_TRUE(peer_.ValidateInvariants());
    }

    TEST_F(Http2PriorityWriteSchedulerTest, UpdateStreamParentToChildNonExclusive)
    {
        /* Tree:
        0
        |
        1
       / \
      2   3
      |
      4
   */
        scheduler_.RegisterStream(1, SpdyStreamPrecedence(0, 100, false));
        scheduler_.RegisterStream(2, SpdyStreamPrecedence(1, 100, false));
        scheduler_.RegisterStream(3, SpdyStreamPrecedence(1, 100, false));
        scheduler_.RegisterStream(4, SpdyStreamPrecedence(2, 100, false));
        scheduler_.UpdateStreamPrecedence(1, SpdyStreamPrecedence(2, 100, false));
        EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(2));
        EXPECT_THAT(scheduler_.GetStreamChildren(1), ElementsAre(3));
        EXPECT_THAT(scheduler_.GetStreamChildren(2), UnorderedElementsAre(1, 4));
        EXPECT_THAT(scheduler_.GetStreamChildren(3), IsEmpty());
        EXPECT_THAT(scheduler_.GetStreamChildren(4), IsEmpty());
        ASSERT_TRUE(peer_.ValidateInvariants());
    }

    TEST_F(Http2PriorityWriteSchedulerTest, UpdateStreamParentToChildExclusive)
    {
        /* Tree:
        0
        |
        1
       / \
      2   3
      |
      4
   */
        scheduler_.RegisterStream(1, SpdyStreamPrecedence(0, 100, false));
        scheduler_.RegisterStream(2, SpdyStreamPrecedence(1, 100, false));
        scheduler_.RegisterStream(3, SpdyStreamPrecedence(1, 100, false));
        scheduler_.RegisterStream(4, SpdyStreamPrecedence(2, 100, false));
        scheduler_.UpdateStreamPrecedence(1, SpdyStreamPrecedence(2, 100, true));
        EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(2));
        EXPECT_THAT(scheduler_.GetStreamChildren(1), UnorderedElementsAre(3, 4));
        EXPECT_THAT(scheduler_.GetStreamChildren(2), ElementsAre(1));
        EXPECT_THAT(scheduler_.GetStreamChildren(3), IsEmpty());
        EXPECT_THAT(scheduler_.GetStreamChildren(4), IsEmpty());
        ASSERT_TRUE(peer_.ValidateInvariants());
    }

    TEST_F(Http2PriorityWriteSchedulerTest,
        UpdateStreamParentToGrandchildNonExclusive)
    {
        /* Tree:
        0
        |
        1
       / \
      2   3
     / \
    4   5
    |
    6
   */
        scheduler_.RegisterStream(1, SpdyStreamPrecedence(0, 100, false));
        scheduler_.RegisterStream(2, SpdyStreamPrecedence(1, 100, false));
        scheduler_.RegisterStream(3, SpdyStreamPrecedence(1, 100, false));
        scheduler_.RegisterStream(4, SpdyStreamPrecedence(2, 100, false));
        scheduler_.RegisterStream(5, SpdyStreamPrecedence(2, 100, false));
        scheduler_.RegisterStream(6, SpdyStreamPrecedence(4, 100, false));
        scheduler_.UpdateStreamPrecedence(1, SpdyStreamPrecedence(4, 100, false));
        EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(4));
        EXPECT_THAT(scheduler_.GetStreamChildren(1), UnorderedElementsAre(2, 3));
        EXPECT_THAT(scheduler_.GetStreamChildren(2), ElementsAre(5));
        EXPECT_THAT(scheduler_.GetStreamChildren(3), IsEmpty());
        EXPECT_THAT(scheduler_.GetStreamChildren(4), UnorderedElementsAre(1, 6));
        EXPECT_THAT(scheduler_.GetStreamChildren(5), IsEmpty());
        EXPECT_THAT(scheduler_.GetStreamChildren(6), IsEmpty());
        ASSERT_TRUE(peer_.ValidateInvariants());
    }

    TEST_F(Http2PriorityWriteSchedulerTest,
        UpdateStreamParentToGrandchildExclusive)
    {
        /* Tree:
        0
        |
        1
       / \
      2   3
     / \
    4   5
    |
    6
   */
        scheduler_.RegisterStream(1, SpdyStreamPrecedence(0, 100, false));
        scheduler_.RegisterStream(2, SpdyStreamPrecedence(1, 100, false));
        scheduler_.RegisterStream(3, SpdyStreamPrecedence(1, 100, false));
        scheduler_.RegisterStream(4, SpdyStreamPrecedence(2, 100, false));
        scheduler_.RegisterStream(5, SpdyStreamPrecedence(2, 100, false));
        scheduler_.RegisterStream(6, SpdyStreamPrecedence(4, 100, false));
        scheduler_.UpdateStreamPrecedence(1, SpdyStreamPrecedence(4, 100, true));
        EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(4));
        EXPECT_THAT(scheduler_.GetStreamChildren(1), UnorderedElementsAre(2, 3, 6));
        EXPECT_THAT(scheduler_.GetStreamChildren(2), ElementsAre(5));
        EXPECT_THAT(scheduler_.GetStreamChildren(3), IsEmpty());
        EXPECT_THAT(scheduler_.GetStreamChildren(4), ElementsAre(1));
        EXPECT_THAT(scheduler_.GetStreamChildren(5), IsEmpty());
        EXPECT_THAT(scheduler_.GetStreamChildren(6), IsEmpty());
        ASSERT_TRUE(peer_.ValidateInvariants());
    }

    TEST_F(Http2PriorityWriteSchedulerTest, UpdateStreamParentToParent)
    {
        scheduler_.RegisterStream(1, SpdyStreamPrecedence(0, 100, false));
        scheduler_.RegisterStream(2, SpdyStreamPrecedence(1, 100, false));
        scheduler_.RegisterStream(3, SpdyStreamPrecedence(1, 100, false));
        for (bool exclusive : { true, false }) {
            scheduler_.UpdateStreamPrecedence(2,
                SpdyStreamPrecedence(1, 100, exclusive));
            EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(1));
            EXPECT_THAT(scheduler_.GetStreamChildren(1), UnorderedElementsAre(2, 3));
            EXPECT_THAT(scheduler_.GetStreamChildren(2), IsEmpty());
            EXPECT_THAT(scheduler_.GetStreamChildren(3), IsEmpty());
        }
        ASSERT_TRUE(peer_.ValidateInvariants());
    }

    TEST_F(Http2PriorityWriteSchedulerTest, UpdateStreamParentToSelf)
    {
        scheduler_.RegisterStream(1, SpdyStreamPrecedence(0, 100, false));
        EXPECT_SPDY_BUG(
            scheduler_.UpdateStreamPrecedence(1, SpdyStreamPrecedence(1, 100, false)),
            "Cannot set stream to be its own parent");
        EXPECT_SPDY_BUG(
            scheduler_.UpdateStreamPrecedence(1, SpdyStreamPrecedence(1, 100, true)),
            "Cannot set stream to be its own parent");
        EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(1));
        EXPECT_THAT(scheduler_.GetStreamChildren(1), IsEmpty());
        ASSERT_TRUE(peer_.ValidateInvariants());
    }

    TEST_F(Http2PriorityWriteSchedulerTest, BlockAndUnblock)
    {
        /* Create the tree.

             0
           / | \
          /  |  \
         1   2   3
        / \   \   \
       4   5   6   7
      /|  / \  |   |\
     8 9 10 11 12 13 14
    / \
   15 16

  */
        scheduler_.RegisterStream(1, SpdyStreamPrecedence(0, 100, false));
        scheduler_.RegisterStream(2, SpdyStreamPrecedence(0, 100, false));
        scheduler_.RegisterStream(3, SpdyStreamPrecedence(0, 100, false));
        scheduler_.RegisterStream(4, SpdyStreamPrecedence(1, 100, false));
        scheduler_.RegisterStream(5, SpdyStreamPrecedence(1, 100, false));
        scheduler_.RegisterStream(8, SpdyStreamPrecedence(4, 100, false));
        scheduler_.RegisterStream(9, SpdyStreamPrecedence(4, 100, false));
        scheduler_.RegisterStream(10, SpdyStreamPrecedence(5, 100, false));
        scheduler_.RegisterStream(11, SpdyStreamPrecedence(5, 100, false));
        scheduler_.RegisterStream(15, SpdyStreamPrecedence(8, 100, false));
        scheduler_.RegisterStream(16, SpdyStreamPrecedence(8, 100, false));
        scheduler_.RegisterStream(12, SpdyStreamPrecedence(2, 100, false));
        scheduler_.RegisterStream(6, SpdyStreamPrecedence(2, 100, true));
        scheduler_.RegisterStream(7, SpdyStreamPrecedence(0, 100, false));
        scheduler_.RegisterStream(13, SpdyStreamPrecedence(7, 100, true));
        scheduler_.RegisterStream(14, SpdyStreamPrecedence(7, 100, false));
        scheduler_.UpdateStreamPrecedence(7, SpdyStreamPrecedence(3, 100, false));
        EXPECT_EQ(0u, scheduler_.GetStreamPrecedence(1).parent_id());
        EXPECT_EQ(0u, scheduler_.GetStreamPrecedence(2).parent_id());
        EXPECT_EQ(0u, scheduler_.GetStreamPrecedence(3).parent_id());
        EXPECT_EQ(1u, scheduler_.GetStreamPrecedence(4).parent_id());
        EXPECT_EQ(1u, scheduler_.GetStreamPrecedence(5).parent_id());
        EXPECT_EQ(2u, scheduler_.GetStreamPrecedence(6).parent_id());
        EXPECT_EQ(3u, scheduler_.GetStreamPrecedence(7).parent_id());
        EXPECT_EQ(4u, scheduler_.GetStreamPrecedence(8).parent_id());
        EXPECT_EQ(4u, scheduler_.GetStreamPrecedence(9).parent_id());
        EXPECT_EQ(5u, scheduler_.GetStreamPrecedence(10).parent_id());
        EXPECT_EQ(5u, scheduler_.GetStreamPrecedence(11).parent_id());
        EXPECT_EQ(6u, scheduler_.GetStreamPrecedence(12).parent_id());
        EXPECT_EQ(7u, scheduler_.GetStreamPrecedence(13).parent_id());
        EXPECT_EQ(7u, scheduler_.GetStreamPrecedence(14).parent_id());
        EXPECT_EQ(8u, scheduler_.GetStreamPrecedence(15).parent_id());
        EXPECT_EQ(8u, scheduler_.GetStreamPrecedence(16).parent_id());
        ASSERT_TRUE(peer_.ValidateInvariants());

        EXPECT_EQ(peer_.TotalChildWeights(0),
            scheduler_.GetStreamPrecedence(1).weight() + scheduler_.GetStreamPrecedence(2).weight() + scheduler_.GetStreamPrecedence(3).weight());
        EXPECT_EQ(peer_.TotalChildWeights(3),
            scheduler_.GetStreamPrecedence(7).weight());
        EXPECT_EQ(peer_.TotalChildWeights(7),
            scheduler_.GetStreamPrecedence(13).weight() + scheduler_.GetStreamPrecedence(14).weight());
        EXPECT_EQ(peer_.TotalChildWeights(13), 0);
        EXPECT_EQ(peer_.TotalChildWeights(14), 0);

        ASSERT_TRUE(peer_.ValidateInvariants());
    }

    TEST_F(Http2PriorityWriteSchedulerTest, HasReadyStreams)
    {
        EXPECT_FALSE(scheduler_.HasReadyStreams());
        scheduler_.RegisterStream(1, SpdyStreamPrecedence(0, 10, false));
        EXPECT_FALSE(scheduler_.HasReadyStreams());
        scheduler_.MarkStreamReady(1, false);
        EXPECT_TRUE(scheduler_.HasReadyStreams());
        scheduler_.MarkStreamNotReady(1);
        EXPECT_FALSE(scheduler_.HasReadyStreams());
        scheduler_.MarkStreamReady(1, true);
        EXPECT_TRUE(scheduler_.HasReadyStreams());
        scheduler_.UnregisterStream(1);
        EXPECT_FALSE(scheduler_.HasReadyStreams());
        ASSERT_TRUE(peer_.ValidateInvariants());
    }

    TEST_F(Http2PriorityWriteSchedulerTest, CalculateRoundedWeights)
    {
        /* Create the tree.

           0
          / \
         1   2
       /| |\  |\
      8 3 4 5 6 7
  */
        scheduler_.RegisterStream(3, SpdyStreamPrecedence(0, 100, false));
        scheduler_.RegisterStream(4, SpdyStreamPrecedence(0, 100, false));
        scheduler_.RegisterStream(5, SpdyStreamPrecedence(0, 100, false));
        scheduler_.RegisterStream(1, SpdyStreamPrecedence(0, 10, true));
        scheduler_.RegisterStream(2, SpdyStreamPrecedence(0, 5, false));
        scheduler_.RegisterStream(6, SpdyStreamPrecedence(2, 1, false));
        scheduler_.RegisterStream(7, SpdyStreamPrecedence(2, 1, false));
        scheduler_.RegisterStream(8, SpdyStreamPrecedence(1, 1, false));

        // Remove higher-level streams.
        scheduler_.UnregisterStream(1);
        scheduler_.UnregisterStream(2);

        // 3.3 rounded down = 3.
        EXPECT_EQ(3, scheduler_.GetStreamPrecedence(3).weight());
        EXPECT_EQ(3, scheduler_.GetStreamPrecedence(4).weight());
        EXPECT_EQ(3, scheduler_.GetStreamPrecedence(5).weight());
        // 2.5 rounded up = 3.
        EXPECT_EQ(3, scheduler_.GetStreamPrecedence(6).weight());
        EXPECT_EQ(3, scheduler_.GetStreamPrecedence(7).weight());
        // 0 is not a valid weight, so round up to 1.
        EXPECT_EQ(1, scheduler_.GetStreamPrecedence(8).weight());
        ASSERT_TRUE(peer_.ValidateInvariants());
    }

    TEST_F(Http2PriorityWriteSchedulerTest, GetLatestEventWithPrecedence)
    {
        EXPECT_SPDY_BUG(scheduler_.RecordStreamEventTime(3, 5),
            "Stream 3 not registered");
        EXPECT_SPDY_BUG(EXPECT_EQ(0, scheduler_.GetLatestEventWithPrecedence(4)),
            "Stream 4 not registered");

        for (int i = 1; i < 5; ++i) {
            int weight = SpdyStreamPrecedence(i).weight();
            scheduler_.RegisterStream(i, SpdyStreamPrecedence(0, weight, false));
        }
        for (int i = 1; i < 5; ++i) {
            EXPECT_EQ(0, scheduler_.GetLatestEventWithPrecedence(i));
        }
        for (int i = 1; i < 5; ++i) {
            scheduler_.RecordStreamEventTime(i, i * 100);
        }
        for (int i = 1; i < 5; ++i) {
            EXPECT_EQ((i - 1) * 100, scheduler_.GetLatestEventWithPrecedence(i));
        }
    }

    // Add ready streams at front and back.
    TEST_F(Http2PriorityWriteSchedulerTest, MarkReadyFrontAndBack)
    {
        scheduler_.RegisterStream(1, SpdyStreamPrecedence(0, 10, false));
        scheduler_.RegisterStream(2, SpdyStreamPrecedence(0, 20, false));
        scheduler_.RegisterStream(3, SpdyStreamPrecedence(0, 20, false));
        scheduler_.RegisterStream(4, SpdyStreamPrecedence(0, 20, false));
        scheduler_.RegisterStream(5, SpdyStreamPrecedence(0, 30, false));

        for (int i = 1; i < 6; ++i) {
            scheduler_.MarkStreamReady(i, false);
        }
        EXPECT_EQ(5u, scheduler_.PopNextReadyStream());
        EXPECT_EQ(2u, scheduler_.PopNextReadyStream());
        scheduler_.MarkStreamReady(2, false);
        EXPECT_EQ(3u, scheduler_.PopNextReadyStream());
        scheduler_.MarkStreamReady(3, false);
        EXPECT_EQ(4u, scheduler_.PopNextReadyStream());
        scheduler_.MarkStreamReady(4, false);
        EXPECT_EQ(2u, scheduler_.PopNextReadyStream());
        scheduler_.MarkStreamReady(2, true);
        EXPECT_EQ(2u, scheduler_.PopNextReadyStream());
        scheduler_.MarkStreamReady(5, false);
        scheduler_.MarkStreamReady(2, true);
        EXPECT_EQ(5u, scheduler_.PopNextReadyStream());
    }

    // Add ready streams at front and back and pop them with
    // PopNextReadyStreamAndPrecedence.
    TEST_F(Http2PriorityWriteSchedulerTest, PopNextReadyStreamAndPrecedence)
    {
        scheduler_.RegisterStream(1, SpdyStreamPrecedence(0, 10, false));
        scheduler_.RegisterStream(2, SpdyStreamPrecedence(0, 20, false));
        scheduler_.RegisterStream(3, SpdyStreamPrecedence(0, 20, false));
        scheduler_.RegisterStream(4, SpdyStreamPrecedence(0, 20, false));
        scheduler_.RegisterStream(5, SpdyStreamPrecedence(0, 30, false));

        for (int i = 1; i < 6; ++i) {
            scheduler_.MarkStreamReady(i, false);
        }
        EXPECT_EQ(std::make_tuple(5u, SpdyStreamPrecedence(0, 30, false)),
            scheduler_.PopNextReadyStreamAndPrecedence());
        EXPECT_EQ(std::make_tuple(2u, SpdyStreamPrecedence(0, 20, false)),
            scheduler_.PopNextReadyStreamAndPrecedence());
        scheduler_.MarkStreamReady(2, false);
        EXPECT_EQ(std::make_tuple(3u, SpdyStreamPrecedence(0, 20, false)),
            scheduler_.PopNextReadyStreamAndPrecedence());
        scheduler_.MarkStreamReady(3, false);
        EXPECT_EQ(std::make_tuple(4u, SpdyStreamPrecedence(0, 20, false)),
            scheduler_.PopNextReadyStreamAndPrecedence());
        scheduler_.MarkStreamReady(4, false);
        EXPECT_EQ(std::make_tuple(2u, SpdyStreamPrecedence(0, 20, false)),
            scheduler_.PopNextReadyStreamAndPrecedence());
        scheduler_.MarkStreamReady(2, true);
        EXPECT_EQ(std::make_tuple(2u, SpdyStreamPrecedence(0, 20, false)),
            scheduler_.PopNextReadyStreamAndPrecedence());
        scheduler_.MarkStreamReady(5, false);
        scheduler_.MarkStreamReady(2, true);
        EXPECT_EQ(std::make_tuple(5u, SpdyStreamPrecedence(0, 30, false)),
            scheduler_.PopNextReadyStreamAndPrecedence());
    }

    class PopNextReadyStreamTest : public Http2PriorityWriteSchedulerTest {
    protected:
        void SetUp() override
        {
            /* Create the tree.

             0
            /|\
           1 2 3
          /| |\
         4 5 6 7
        /
       8

    */
            scheduler_.RegisterStream(1, SpdyStreamPrecedence(0, 100, false));
            scheduler_.RegisterStream(2, SpdyStreamPrecedence(0, 100, false));
            scheduler_.RegisterStream(3, SpdyStreamPrecedence(0, 100, false));
            scheduler_.RegisterStream(4, SpdyStreamPrecedence(1, 100, false));
            scheduler_.RegisterStream(5, SpdyStreamPrecedence(1, 100, false));
            scheduler_.RegisterStream(6, SpdyStreamPrecedence(2, 100, false));
            scheduler_.RegisterStream(7, SpdyStreamPrecedence(2, 100, false));
            scheduler_.RegisterStream(8, SpdyStreamPrecedence(4, 100, false));

            // Set all nodes ready to write.
            for (SpdyStreamId id = 1; id <= 8; ++id) {
                scheduler_.MarkStreamReady(id, false);
            }
        }

        AssertionResult PopNextReturnsCycle(
            std::initializer_list<SpdyStreamId> stream_ids)
        {
            int count = 0;
            const int kNumCyclesToCheck = 2;
            for (int i = 0; i < kNumCyclesToCheck; i++) {
                for (SpdyStreamId expected_id : stream_ids) {
                    SpdyStreamId next_id = scheduler_.PopNextReadyStream();
                    scheduler_.MarkStreamReady(next_id, false);
                    if (next_id != expected_id) {
                        return AssertionFailure() << "Pick " << count << ": expected stream "
                                                  << expected_id << " instead of " << next_id;
                    }
                    if (!peer_.ValidateInvariants()) {
                        return AssertionFailure() << "ValidateInvariants failed";
                    }
                    ++count;
                }
            }
            return AssertionSuccess();
        }
    };

    // When all streams are schedulable, only top-level streams should be returned.
    TEST_F(PopNextReadyStreamTest, NoneBlocked)
    {
        EXPECT_TRUE(PopNextReturnsCycle({ 1, 2, 3 }));
    }

    // When a parent stream is blocked, its children should be scheduled, if
    // priorities allow.
    TEST_F(PopNextReadyStreamTest, SingleStreamBlocked)
    {
        scheduler_.MarkStreamNotReady(1);

        // Round-robin only across 2 and 3, since children of 1 have lower priority.
        EXPECT_TRUE(PopNextReturnsCycle({ 2, 3 }));

        // Make children of 1 have equal priority as 2 and 3, after which they should
        // be returned as well.
        scheduler_.UpdateStreamPrecedence(1, SpdyStreamPrecedence(0, 200, false));
        EXPECT_TRUE(PopNextReturnsCycle({ 4, 5, 2, 3 }));
    }

    // Block multiple levels of streams.
    TEST_F(PopNextReadyStreamTest, MultiLevelBlocked)
    {
        for (SpdyStreamId stream_id : { 1, 4, 5 }) {
            scheduler_.MarkStreamNotReady(stream_id);
        }
        // Round-robin only across 2 and 3, since children of 1 have lower priority.
        EXPECT_TRUE(PopNextReturnsCycle({ 2, 3 }));

        // Make 8 have equal priority as 2 and 3.
        scheduler_.UpdateStreamPrecedence(1, SpdyStreamPrecedence(0, 200, false));
        EXPECT_TRUE(PopNextReturnsCycle({ 8, 2, 3 }));
    }

    // A removed stream shouldn't be scheduled.
    TEST_F(PopNextReadyStreamTest, RemoveStream)
    {
        scheduler_.UnregisterStream(1);

        // Round-robin only across 2 and 3, since previous children of 1 have lower
        // priority (the weight of 4 and 5 is scaled down when they are elevated to
        // siblings of 2 and 3).
        EXPECT_TRUE(PopNextReturnsCycle({ 2, 3 }));

        // Make previous children of 1 have equal priority as 2 and 3.
        scheduler_.UpdateStreamPrecedence(4, SpdyStreamPrecedence(0, 100, false));
        scheduler_.UpdateStreamPrecedence(5, SpdyStreamPrecedence(0, 100, false));
        EXPECT_TRUE(PopNextReturnsCycle({ 4, 5, 2, 3 }));
    }

    // Block an entire subtree.
    TEST_F(PopNextReadyStreamTest, SubtreeBlocked)
    {
        for (SpdyStreamId stream_id : { 1, 4, 5, 8 }) {
            scheduler_.MarkStreamNotReady(stream_id);
        }
        EXPECT_TRUE(PopNextReturnsCycle({ 2, 3 }));
    }

    // If all parent streams are blocked, children should be returned.
    TEST_F(PopNextReadyStreamTest, ParentsBlocked)
    {
        for (SpdyStreamId stream_id : { 1, 2, 3 }) {
            scheduler_.MarkStreamNotReady(stream_id);
        }
        EXPECT_TRUE(PopNextReturnsCycle({ 4, 5, 6, 7 }));
    }

    // Unblocking streams should make them schedulable.
    TEST_F(PopNextReadyStreamTest, BlockAndUnblock)
    {
        EXPECT_TRUE(PopNextReturnsCycle({ 1, 2, 3 }));
        scheduler_.MarkStreamNotReady(2);
        EXPECT_TRUE(PopNextReturnsCycle({ 1, 3 }));
        scheduler_.MarkStreamReady(2, false);
        // Cycle order permuted since 2 effectively appended at tail.
        EXPECT_TRUE(PopNextReturnsCycle({ 1, 3, 2 }));
    }

    // Block nodes in multiple subtrees.
    TEST_F(PopNextReadyStreamTest, ScatteredBlocked)
    {
        for (SpdyStreamId stream_id : { 1, 2, 6, 7 }) {
            scheduler_.MarkStreamNotReady(stream_id);
        }
        // Only 3 returned, since of remaining streams it has highest priority.
        EXPECT_TRUE(PopNextReturnsCycle({ 3 }));

        // Make children of 1 have priority equal to 3.
        scheduler_.UpdateStreamPrecedence(1, SpdyStreamPrecedence(0, 200, false));
        EXPECT_TRUE(PopNextReturnsCycle({ 4, 5, 3 }));

        // When 4 is blocked, its child 8 should take its place, since it has same
        // priority.
        scheduler_.MarkStreamNotReady(4);
        EXPECT_TRUE(PopNextReturnsCycle({ 8, 5, 3 }));
    }

} // namespace test
} // namespace net
